function [theta,stdDKj,stdW,CovDKj,yhat,R2,CovW,CovDK,stdDK] = HszDk5cPs(y,x,z,yhatQ,m,ScaleByNtQ)
% Note: changed the order of stdDKj (CovDKj) and stdDK (CovDK), Benjamin Born on 27072018

%HszDk5cPs  LS and Driscoll-Kray standard errors for unbalanced panel, assuming
%           x are the same across individuals, while z are time-varying individual
%           characteristics. The effective regressors are kron(z,x).
%
%  Usage: [theta,stdDK,stdW,CovDK,yhat,R2,CovW,CovDKj,stdDKj] = HszDk5cPs(y,x,z,yhatQ,m);
%
%  Input:    y             TxN matrix with the dependent variable, y(t,i) is for period t, individual i
%            x             TxK matrix with K factors that are common for all investors
%            z             TxNxL matrix with L (time-varying) individual characteristics
%            yhatQ         (optional) scalar, 1: generate and report fitted values
%            m             (optional), scalar, number of lags in covariance estimation
%            ScaleByNtQ    (optional), scalar, 1: scales all moment conditions by N(t), not by N, [0],
%                          can be used to replicate a portfolio (Calendar time) approach
%
%  Output:   theta         (K*L)x1 vector, LS estimates of regression coeefficients on kron(z,x)
%            stdDK         (K*L)x1 vector, Driscoll-Kraay standard errors
%            stdW          (K*L)x1 vector, White's standard errors
%            CovDK         (K*L)x(K*L) matrix, Driscoll-Kraay covariance matrix
%            yhat          TxN matrix with fitted values
%            R2            scalar, (pseudo-) R2
%            CovW          covariance matrix, White's
%            CovDKj        covariance matrix, DK with lags
%            stdDKj        standard errors, DK with lags
%
%
%  Notice:   (a) the effective regressors are kron(z,x). For instance, with z = [1,z1] and
%                x = [1,x1,x2,x3], we have [1,x1,x2,x3,z1,z1*x1,z1*x2,z1*x3].
%
%            (b) y,x,z can be single precision, since they are transformed into double precision
%                inside the function
%
%  Paul.Soderlind@unisg.ch   May 2010, Dec 2010, Jun 2011, Dec 2011
%  Benjamin Born and Johannes Pfeifer, 2013-2015
%------------------------------------------------------------------------------

if (nargin < 4) || isempty(yhatQ);
    yhatQ = 0;
end;
if (nargin < 5) || isempty(m);
    m = 0;
end;
if (nargin < 6) || isempty(ScaleByNtQ);
    ScaleByNtQ = 0;
end;


[T,N] = size(y);
L     = size(z,3);
KL    = size(x,2)*L;


xx = 0;                              %Sum[x(t)*x(t)',t=1:T]
xy = 0;                              %Sum[x(t)*y(t),t=1:T]
Nb = NaN(T,1);                       %effective number of obs, after pruning NaNs
for t = 1:T;                           %loop over time
    y_t   = double(y(t,:)');                %dependent variable, Nx1
    x0_t  = repmat(double(x(t,:)),N,1);     %factors, NxK
    z_t   = reshape(double(z(t,:,:)),N,L);  %NxL, better than squeeze (cf 2-d arrays)
    x_t   = HDirProdPs(z_t,x0_t);           %effective regressors, z_t is NxL, x_t is NxK
    yx_t  = delete_NaN_countries([y_t,x_t]);              %pruning NaNs
    if ScaleByNtQ == 1;
        Nb(t) = size(yx_t,1);
    else
        Nb(t) = N;
    end;
    if ~isempty(yx_t);                   %don't accumulate [] to xx and xy (generates [])
        y_t = yx_t(:,1);
        x_t = yx_t(:,2:end);
        xx  = xx + x_t'*x_t/Nb(t);
        xy  = xy + x_t'*y_t/Nb(t);
    else
        fprintf('\nNo variables left at time %d\n',t)
    end;
end;
clear yx_t;

Tb  = sum(Nb>0);                    %number of effective time periods
xx  = xx/Tb;
xy  = xy/Tb;
theta  = xx\xy;                      %ols estimates, solves xx*theta = xy


if yhatQ == 1;
    yhat = NaN(T,N);
else
    yhat = [];
end;
omega0DK = zeros(KL,KL);               %DK, lag 0
omega0W  = omega0DK;                   %White's
omegajDK = zeros([KL,KL,m]);           %DK, lags 1 to m
h_tLag   = zeros(m,KL);                %lag1;lag2;...,lagm
for t = 1:T;                           %loop over time
    y_t    = double(y(t,:)');
    x0_t   = repmat(double(x(t,:)),N,1);
    z_t    = reshape(double(z(t,:,:)),N,L);
    x_t    = HDirProdPs(z_t,x0_t);
    yhat_t = x_t*theta;
    r_t    = y_t - yhat_t;
    rx_t   = delete_NaN_countries([r_t,x_t]);
    if ~isempty(rx_t);                   %don't accumulate [] to omega0DK
        r_t    = rx_t(:,1);
        x_t    = rx_t(:,2:end);
        hi_t   = x_t.*repmat(r_t,1,KL);    %moment condition for (i,t)
        h_t    = sum(hi_t,1)/Nb(t);
        omega0DK = omega0DK + h_t'*h_t;
        omega0W  = omega0W + hi_t'*hi_t/Nb(t)^2;
        for j = 1:m;
            omegajDK(:,:,j) = omegajDK(:,:,j) + h_t'*h_tLag(j,:);    %h(t)*h(t-j)'
        end;
        h_tLag = [h_t;h_tLag(1:end-1,:)];  %update only if ~isempty(rx_t), effectively disregarding t if no data
    end;
    if yhatQ == 1;
        yhat(t,:) = yhat_t';
    end;
end;
Shat  = omega0DK/Tb^2;                   %estimate of S, DK
Shatw = omega0W/Tb^2;                    %estimate of S, White's
Shatj = omega0DK/Tb^2;
for j = 1:m;
    Shatj = Shatj + (1-j/(m+1))*(omegajDK(:,:,j)+omegajDK(:,:,j)')/Tb^2 ;
end;


zx_1  = inv(xx);
CovDK = zx_1 * Shat * zx_1';                      %covariance matrix, DK
stdDK = sqrt( diag(CovDK) );                      %standard errors, DK
CovW  = zx_1 * Shatw * zx_1';                     %covariance matrix, White's
stdW  = sqrt( diag(CovW) );                       %standard errors, White's
CovDKj = zx_1 * Shatj * zx_1';                    %covariance matrix, DK with lags
stdDKj = sqrt( diag(CovDKj) );                    %standard errors, DK with lags


if yhatQ == 1;
    yy   = delete_NaN_countries([yhat(:),y(:)]);
    R2   = corrcoef(yy);
    R2   = R2(1,2)^2;
    clear yy;
else
    R2   = [];
end;
%------------------------------------------------------------------------------
end
function x=delete_NaN_countries(x)

miss=isnan(x);
[m,n]=size(x);
if n==1,
    x(miss) = [];
else
    x(any(miss'),:) = []; %delete countries (in rows) where full set of regressors is not present (a NaN in the column)
end;

end

function z = HDirProdPs(x,y)
%HorizDirectProdPs    Calculates horizontal direct product of two matrices with equal number of rows.
%                     z(i,:) is the Kronecker product of x(i,:) and y(i,:).
%
%  Usage:    z = HDirProdPs(x,y);
%
%  Input:    x      T x Kx matrix
%            y      T x Ky matrix
%
%  Output:   z      T x (Kx*Ky)
%
%  Example:   x = [ 1 2;
%                   3 4];
%             y = [5 6 1;
%                  7 8 1];
%
%             the z = HDirProdPs(x,y) gives
%
%             z = [ 5  6  1  10  12  2;
%                  21 24  3  28  32  4];
%
%  Remark: (a)  This could also (but much more slowly) be calculated as
%               z = [];
%               for i=1:Ky;
%                 z = [z,x .* repmat(y(i,:),1,Ky)];
%               end;
%               or
%               for i=1:T;
%                 z(i,:) = kron(x(i,:),y(i,:));
%               end;
%
%           (b) This is the same as x *~ y in Gauss
%
% Paul.Soderlind@unisg.ch, Feb 2010
%----------------------------------------------------------------------------*/

[T,Kx] = size(x);         %rows and columns in x
Ky     = size(y,2);       %columns in y

vvx = kron(1:Kx,ones(1,Ky));
vvy = kron(ones(1,Kx),1:Ky);

z = x(:,vvx) .* y(:,vvy);

% z = zeros(T,Kx*Ky);          %an alternative
% z(:,:) = NaN;
% for i=1:Kx;                  %loop over columns in x
%   vv = ((i-1)*Ky+1):(i*Ky);
%   z(:,vv) = repmat(x(:,i),1,Ky).*y;
% end;

% z = zeros(T,Kx*Ky);          %another alternative
% z(:,:) = NaN;
% for i=1:T;                    %loop over rows
%   z(i,:) = kron(x(i,:),y(i,:));
% end;
%----------------------------------------------------------------------------*/
end
